#!/usr/bin/env python

# Raja S
# Dec 2008
# GPL

"""Test the qrs detection accuracy using database ECGs and comparing
against their annotations"""

import os
import scipy
import glob
import time
import commands
import pywfdb
from pt_qrsdetect import Ecg


#-----------------------------
sampledir = '/data/tmp/mit_db'
annotator = 'atr'
#----------------------------

def qrsdetect_singlelead(record):
    """Detect qrs peaks in first lead of record"""
    recname = record.name
    signal = record.read(record.signal_names[0])
    ecg = Ecg(signal, {'samplingrate' : record.frequency})
    ecg.qrsDetect()
#    ecg.write_ann(os.path.join(sampledir, record.name + '.slr'))
    ecg._write_ann(recname + '.ptx')

    ret, msg = commands.getstatusoutput("cat %s | wrann -r %s -a ptb" %(recname+'.ptx', recname))
    assert ret == 0, 'wrann failed, ' + msg

    ret, msg = commands.getstatusoutput("bxb -r %s -a atr ptb -l res.txt res2.txt" %(recname))
    assert ret == 0, 'bxb failed, ' + msg
                       
def qrsdetect_multilead(record):
    """Detect peaks using all leads"""
    signal = []
    for lead in record.signal_names:
        signal.append(record.read(lead))

    ecg = Ecg(scipy.array(signal).transpose(),
              {'samplingrate' : record.frequency})
    ecg.qrs_detect_multiple_leads()
    ecg.write_ann(os.path.join(sampledir, record.name + '.mlr'))
    
def main():
    """Run qrs detection for all records in directory
    containing samples using single lead and using
    multiple leads"""
    # make list of records in sample directory
    os.chdir(sampledir)
    datfiles = glob.glob('*.dat')
    records = [datfilename.rstrip('.dat') for datfilename
               in datfiles]
    assert len(records) > 0, "No '.dat' files found in %s" %(sampledir)
    print "Found %s records in sample directory" %(len(records)) 

    # confirm that header and annotation files are present for each record
    for headerfile in [recordname+'.hea' for recordname
                       in records]:
        assert os.path.isfile(headerfile),\
               'Header file %s missing' %(headerfile)

    for annotfile in [recordname+'.'+annotator for recordname
                      in records]:
        assert os.path.isfile(annotfile),\
               'Annotation file %s missing' %(annotfile)

    # detect qrs
    for record in records:
        stime = time.time()
        print "------------"
        print "starting record ", record
        rec = pywfdb.Record(os.path.join(sampledir, record))
        qrsdetect_singlelead(rec)
#        qrsdetect_multilead(rec)
        print "completed detection for record %s in %s seconds\n" %(record, time.time()-stime)

def test():
    os.chdir('/data/tmp/mit_db')
    record = '100'
    rec = pywfdb.Record(os.path.join(sampledir, record))
    signal = rec.read(rec.signal_names[0])
    ecg = Ecg(signal, {'samplingrate' : rec.frequency})
    ecg.qrsDetect()
    ecg._write_ann(record + '.ptx')

    ret, msg = commands.getstatusoutput("cat %s | wrann -r %s -a ptb" %(record+'.ptx', record))
    assert ret == 0, 'wrann failed, ' + msg

    ret, msg = commands.getstatusoutput("bxb -r %s -a atr ptb" %(record))
    assert ret == 0, 'bxb failed, ' + msg
    print msg    
    
#    ecg.visualize_qrs_detection()
#    ecg.visualize_qrs_detection2()


if __name__ == '__main__':
    main()
